
import PyRakNet
import struct
import pickle

# TODO: Think about what to do with all of the RakNet built-in packet types. They
#        need to be handled somewhere, but the handling will likely be quite
#        game-specific. Probably some sort of integration with the engine's event
#        system would be the way to do it. In fact this could be a recurring issue
#        with other engine components, so a general solution would be nice. The
#        main thing is to avoid coupling as much as possible, unless we just
#        accept that the event system needs to be coupled to everything. But
#        perhaps something like being able to bind the packet events to arbitrary
#        functions, which could be handler or intermediary functions elsewhere
#        in the engine where the coupling would take place.

# TODO: Check server as client!

class Network(object):
    """Handles all of the networking for the engine

       Dependencies: pyraknet, struct, pickle
    """
    # Node types
    NET_UNDEFINED = 0
    NET_SERVER = 1
    NET_CLIENT = 2
    
    # Our packet ids
    ID_RPC = PyRakNet.ID_USER_PACKET_ENUM

    # Attribute synchronisation types
    SYNCH_ON_REPLICATE_ONLY = 0
    SYNCH_CONTINUOUSLY = 1
    
    def __init__(self):
        self.node = None
        self.node_type = Network.NET_UNDEFINED
        self.host = None
        self.thread_sleep_timer = 30
        # The registry of remote procedures
        self.rpcs = {}
        # The registry of remote objects and types
        self.remote_objects = {}
        self.free_object_keys = []
        #self.remote_types = {} # Don't need this as both types and instances are stored in remote_objects

    def startAsServer(self, max_clients=32, port=1234):
        if self.node_type != Network.NET_UNDEFINED:
            print "Shutting down node"
            self.node.Shutdown(0)
        print "Starting up server..."
        self.node_type = Network.NET_SERVER
        self.node = PyRakNet.RakPeer()
        self.node.Startup(max_clients, self.thread_sleep_timer, PyRakNet.SocketDescriptor(port, ""))
        self.node.SetMaximumIncomingConnections(max_clients)
        self.node.SetOccasionalPing(True)
        print 'Server initialised, waiting for connections...'
        
    def startAsClient(self, server_ip='localhost', server_port=1234):
        if self.node_type != Network.NET_UNDEFINED:
            print "Shutting down node"
            self.node.Shutdown(0)
        print "Initialising client..."
        self.node_type = Network.NET_CLIENT
        self.node = PyRakNet.RakPeer()
        self.node.Startup(1, self.thread_sleep_timer, PyRakNet.SocketDescriptor(0, ""))
        self.node.Connect(server_ip, server_port, "")
        self.node.SetOccasionalPing(True)
        print 'Client initialised, connecting to server...'
    
    def update(self):
        self.handleIncomingPackets()
        self.sendUpdates()
    
    def handleIncomingPackets(self):
        packet = self.node.Receive()
        if packet:
            type_id = ord(packet.data[0])
            # Check for timestamped packets
            if type_id == PyRakNet.ID_TIMESTAMP:
                # Unpack the timestamp
                # Note: The '5' used in here is the length of a packed timestamp, which is a uint
                # Note: struct.unpack returns a tuple even when only 1 value is unpacked
                time = struct.unpack('<I', packet.data[1:5])[0]
                # Then read the actual packet type identifier
                type_id = ord(packet.data[5])
                # And extract the payload
                data = packet.data[6:]
                
            else:
                # No timestamp
                time = None
                # Extract the payload
                data = packet.data[1:]
            
            if self.node_type == Network.NET_SERVER:
            
                # We've just received a packet from a client
                if type_id == Network.ID_RPC:
                    print "Server received RPC from ", packet.address, ", forwarding to other clients."
                    # Forward the packet to all of the other clients
                    self.send( packet.data, \
                               PyRakNet.PacketPriority.MEDIUM_PRIORITY, \
                               PyRakNet.PacketReliability.RELIABLE, 0, \
                               packet.address, True)
                elif type_id == PyRakNet.ID_NEW_INCOMING_CONNECTION:
                    """A remote system has successfully connected."""
                    
                    print "New incoming connection." # Replicating existing network objects that I have authority over.
                    
                    # New connections must be given all existing network objects
                    for remote_object in self.remote_objects.values():
                        # Check that it's an remote instance object and not a remote type
                        if hasattr(remote_object, '_network_key_'):
                            print "  Replicating remote object with network key ", remote_object._network_key_, " and args ", remote_object._network_init_args_, remote_object._network_init_kwargs_
                            
                            # TODO: This is not really very nice. Ideally we could use the RPC system as is.
                            #       The problem is that it doesn't allow us to send to a specific client.
                            
                            # Create an RPC packet to represent the remote replication of the object
                            new_data = self.__buildRPCPacket(False, \
                                                             remote_object._create_remotely_._network_key_, \
                                                             remote_object.__class__, \
                                                             remote_object._network_key_, \
                                                             *remote_object._network_init_args_, \
                                                             **remote_object._network_init_kwargs_)
                            # Send it to the new client
                            self.send( new_data, \
                                       PyRakNet.PacketPriority.MEDIUM_PRIORITY, \
                                       PyRakNet.PacketReliability.RELIABLE_ORDERED, 0, \
                                       packet.address, False)
                            """
                            # Create the object locally
                            new_instance = remote_object._create_remotely_(remote_object.__class__, \
                                                                           remote_object._network_key_, \
                                                                           *remote_object._network_init_args_, \
                                                                           **remote_object._network_init_kwargs_)
                            """
                            # Call its __init__ method locally with the saved arguments
                            #new_instance.__init__(*remote_object._network_init_args_, **remote_object._network_init_kwargs_)
                            # TODO: Now synch all of its SYNCH_ON_REPLICATE_ONLY attributes
                elif type_id == PyRakNet.ID_DISCONNECTION_NOTIFICATION:
                    """The system specified in Packet::systemAddress has disconnected from us. For the client, this would mean the server has shutdown."""
                    pass
                elif type_id == PyRakNet.ID_CONNECTION_LOST:
                    """Reliable packets cannot be delivered to the system specifed in Packet::systemAddress. The connection to that system has been closed."""
                    pass
                elif type_id == PyRakNet.ID_MODIFIED_PACKET:
                    """A packet has been tampered with in transit. The sender is contained in Packet::systemAddress."""
                    pass
                else:
                    print "Received unhandled packet with type_id: ", type_id
                
            elif self.node_type == Network.NET_CLIENT:
            
                # We've just received a packet from the server
                if type_id == Network.ID_RPC:
                    # Unpack the RPC packet data
                    procedure_key, args, kwargs, latency = self.__unpackRPCPacket(data, time)
                    # Look up the procedure in our rpc registry
                    procedure = self.rpcs[procedure_key]
                    if latency != None:
                        # Set the latency member variable of the function
                        procedure.network_latency = latency
                    print "Received RPC with args: ", args, kwargs
                    # Call the procedure locally
                    procedure(*args, **kwargs)
                    
                    # TODO: Think about how and where to deal with functions being
                    #        called authoritatively by the server when they may
                    #        already have been called locally by the client. Ideally
                    #        this code would be in a layer on top of the low-level
                    #        network code in here.
                    #       The code for this could go in the procedure itself.
                    #        e.g., createEntity(): if entity already exists: ...

                elif type_id == PyRakNet.ID_CONNECTION_REQUEST_ACCEPTED:
                    """In a client/server environment, our connection request to the server has been accepted."""
                    pass
                elif type_id == PyRakNet.ID_CONNECTION_ATTEMPT_FAILED:
                    """Sent to the player when a connection request cannot be completed due to inability to connect."""
                    pass
                elif type_id == PyRakNet.ID_ALREADY_CONNECTED:
                    """Sent when we requested a connection with a system we are already connected to."""
                    pass
                elif type_id == PyRakNet.ID_NO_FREE_INCOMING_CONNECTIONS:
                    """The system we attempted to connect to is not accepting new connections."""
                    pass
                elif type_id == PyRakNet.ID_DISCONNECTION_NOTIFICATION:
                    """The system specified in Packet::systemAddress has disconnected from us. For the client, this would mean the server has shutdown."""
                    pass
                elif type_id == PyRakNet.ID_CONNECTION_LOST:
                    """Reliable packets cannot be delivered to the system specifed in Packet::systemAddress. The connection to that system has been closed."""
                    pass
                elif type_id == PyRakNet.ID_CONNECTION_BANNED:
                    """We are banned from the system we attempted to connect to."""
                    pass
                elif type_id == PyRakNet.ID_INVALID_PASSWORD:
                    """The remote system is using a password and has refused our connection because we did not set the correct password."""
                    pass
                elif type_id == PyRakNet.ID_REMOTE_DISCONNECTION_NOTIFICATION:
                    """In a client/server environment, a client other than ourselves has disconnected gracefully. Packet::systemAddress is modified to reflect the systemAddress of this client."""
                    pass
                elif type_id == PyRakNet.ID_REMOTE_CONNECTION_LOST:
                    """In a client/server environment, a client other than ourselves has been forcefully dropped. Packet::systemAddress is modified to reflect the systemAddress of this client."""
                    pass
                elif type_id == PyRakNet.ID_REMOTE_NEW_INCOMING_CONNECTION:
                    """In a client/server environment, a client other than ourselves has connected. Packet::systemAddress is modified to reflect the systemAddress of this client."""

                    print "New incoming connection. Replicating existing network objects that I have authority over."
                    
                    # New connections must be given all existing network objects
                    for remote_object in self.remote_objects:
                        # Check that it's an remote instance object and not a remote type
                        if hasattr(remote_object, '_network_key_'):
                            # TODO: Work out how to get the set of arguments to send to the new object's __init__
                            #remote_object._remote_replicate_(remote_object._network_key_, (?,))
                            pass
                else:
                    print "Received unknown\unhandled packet with type_id: ", type_id
                    
            else:
                raise Exception("Network error: network object has not been initialised as a client or server.")

    def sendUpdates(self):
        for remote_object in self.remote_objects:
            # Check that it's an remote instance object and not a remote type
            # TODO: It might be an idea to optimise this by storing network types in a separate list
            if hasattr(remote_object, '_network_key_'):
                for remote_attr in remote_object._remote_attributes_:
                    pass
                    #if remote_attr.

    def __timestampPacket(self, data):
        timestamp = struct.pack( "<I", PyRakNet.GetTime() )
        return chr(PyRakNet.ID_TIMESTAMP) + timestamp + data
    
    def __buildRPCPacket(self, is_timestamped, procedure_key, *args, **kwargs):
        # Replace references to network objects with the objects' network keys
        # Note: netIDs are packed inside a tuple so that they can be identified
        newargs = []
        for arg in args:
            if type(arg) == tuple: newargs.append((arg,))
            elif hasattr(arg,'_network_key_'): newargs.append((arg._network_key_,))
            # TODO: Could consider wrapping remote types in a list instead of a tuple and using two seperate arrays for lookup
            elif hasattr(arg,'_network_type_key_'): newargs.append((arg._network_type_key_,))
            else: newargs.append(arg)
        newkwargs = {}
        for key, value in kwargs.items():
            if type(value) == tuple: newkwargs[key] = (value,)
            elif hasattr(value,'_network_key_'): newkwargs[key] = (value._network_key_,)
            # TODO: Could consider wrapping remote types in a list instead of a tuple and using two seperate arrays for lookup
            elif hasattr(value,'_network_type_key_'): newkwargs[key] = (value._network_type_key_,)
            else: newkwargs[key] = value
        
        argdata = pickle.dumps((procedure_key, newargs, newkwargs), protocol=2)
        
        data = chr(Network.ID_RPC) + argdata
        if is_timestamped:
            data = self.__timestampPacket(data)
        return data

    def __unpackRPCPacket(self, data, timestamp):
        # Unpack the RPC packet data
        procedure_key, args, kwargs = pickle.loads(data)
        
        # Replace encoded entity ids with their actual instances
        newargs=[]
        for arg in args:
            # Check for network keys wrapped in tuples
            if type(arg) == tuple:
                # Check for anything but another tuple inside the tuple
                if type(arg[0]) != tuple:
                    # We found an encoded network key
                    # Check if we know this object
                    if arg[0] in self.remote_objects:
                        newargs.append(self.remote_objects[arg[0]])
                    else:
                        # Oops, someone buggered up
                        raise Exception("RPC '" + procedure_key + "'s argument '" + str(arg[0]) + "' refers to unknown remote object")
                # Must be a plain old tuple we encoded by putting it inside another one
                else: newargs.append(arg[0])
            # Just a plain old arg
            else: newargs.append(arg)
        args = newargs
        newkwargs={}
        for key, value in kwargs.items():
            # Check for network keys wrapped in tuples
            if type(value) == tuple:
                # Check for anything but another tuple inside the tuple
                if type(value[0]) != tuple:
                    # We found an encoded network key
                    # Check if we know this object
                    if value[0] in self.remote_objects:
                        newkwargs[key] = self.remote_objects[value[0]]
                    else:
                        # Oops, someone buggered up
                        raise Exception("RPC '" + procedure_key + "'s keyword argument '" + str(key) + "' with value '" + str(value) + "' refers to unknown remote object")
                # Must be a plain old tuple we encoded by putting it inside another one
                else: newkwargs[key] = value[0]
            # Just a plain old arg
            else: newkwargs[key] = value
        kwargs = newkwargs

        if timestamp == None:
            latency = None
        else:
            latency = PyRakNet.GetTime() - timestamp
            #args.append(latency)
        
        return procedure_key, args, kwargs, latency

    # Note: Could add args to the decorator to specify priority and reliability
    def remote_procedure(self, authority, is_timestamped=False):
        """Produces decorators for making functions call over the network
           
           Note: timestamped RPCs have a network_latency attribute added to them
        """
        def decorator(func):
            def replacement(*args, **kwargs):
                if is_timestamped:
                    # Add a network latency member variable to the function
                    replacement.network_latency = 0
                # If we're the authority over this procedure
                if self.node_type == authority:
                    # Pack up the function and send it to the server
                    data = self.__buildRPCPacket(is_timestamped, replacement._network_key_, *args, **kwargs)
                    # Send to server\all clients
                    self.send( data, \
                               PyRakNet.PacketPriority.MEDIUM_PRIORITY, \
                               PyRakNet.PacketReliability.RELIABLE_ORDERED, 0 )
                ## Else if we're the server and it's a client-authority procedure
                #elif self.node_type == NET_SERVER:
                #    # Pack up the function and send it to all of the other clients
                #    data = self.__buildRPCPacket(replacement._key_, *args, **kwargs)
                #    self.send(data, PacketPriority.MEDIUM_PRIORITY, PacketReliability.RELIABLE, 0, )
                
                func(*args, **kwargs)
                
            # TODO: Double check that this generated key will be unique
            #       It will not be! How the hell can I get the class that the function is defined in?
            # Check if the function is a member function
            if hasattr(func, 'im_class'):
            #if hasattr(func, '_network_key_'):
                key = str(func.im_class) + '.' + str(func.__name__)
                print "Registering member RPC", func, "with key", key
            else:
                key = str(func.__module__) + '.' + str(func.__name__)
                print "Registering non-member RPC", func, "with key", key
            self.rpcs[key] = func
            replacement._network_key_ = key
            return replacement
        return decorator

    # TODO: This is not used correctly! Need to set sumvar = newObjectKey().next    
    def newObjectKey(self):
        """Obtains a new remote object key
        """
        i = -1
        while True:
            if len(self.free_object_keys) == 0:
                i += 1
                yield i
            else:
                i = self.free_object_keys.pop()
                yield i
    
    def freeObjectKey(self, value):
        """Frees a remote object key for reuse
        """
        self.free_object_keys.append(value)

    def REMOTE_OBJECT(network_self, authority, is_timestamped=False):
        """Produces decorators for making a classes into remote objects

           authority:      The authority over creation and destruction of this object
           is_timestamped: Whether to timestamp creation of this object
           Note: Timestamped objects must accept the latency as the last argument of their __init__ method
           TODO: Is this going to be bad if we want an __init__ which accepts *args and **kwargs?

           The decorator modifies the __init__ method to automatically create
            instances of the object across the network. It also adds the
            self-explanatory removeFromNetwork() method.
        """
        def CLASS_DECORATOR(Original):
            
            class Replacement(Original):

                def __new__(cls, *args, **kwargs):
                    print cls, "__new__", args, kwargs
                    # Instantiate the object
                    new_instance = super(Original, cls).__new__(cls, *args, **kwargs)
                    
                    # Check that a super class hasn't already set the network key
                    if not hasattr(new_instance, '_network_key_'):
                        if network_self.node_type == authority:
                            # Generate a new remote object key
                            new_instance._network_key_ = network_self.newObjectKey().next()
                            
                            Replacement._create_remotely_(cls, new_instance._network_key_, *args, **kwargs)
                        else: # If non-auth
                            # Check that the call came from the remote authority
                            # Note: This is very inelegant, but what can you do.
                            if len(args) == 0 and len(kwargs) == 1 and 'network_key' in kwargs:
                                # Set its network key to the one we've been given by the authority
                                new_instance._network_key_ = kwargs['network_key']
                            else:
                                raise Exception("Attempt to instantiate remote object '" + str(cls) + "' by non-authority.")
                                return None

                    # Add the remote object to the network's registry
                    network_self.remote_objects[new_instance._network_key_] = new_instance
                    return new_instance

                @staticmethod
                @network_self.remote_procedure(authority, is_timestamped=False)
                def _create_remotely_(cls, _network_key, *args, **kwargs):
                    if network_self.node_type != authority:
                        new_instance = cls.__new__(cls, network_key=_network_key)
                        print "Remote object created on non-authority with network key: ", _network_key, " and args ", args, kwargs
                        new_instance.__init__(*args, **kwargs)
                        # Note: This return value is only useful on the remote machine.
                        return new_instance
                
                def __init__(self, *args, **kwargs):
                    print self.__class__.__name__, "__init__", args, kwargs
                    
                    # Save the __init__ args so that they can be given later if the object needs to be replicated
                    self._network_init_args_ = args
                    self._network_init_kwargs_ = kwargs

                    # Add dicts for synchronised attributes
                    self._synched_on_replicate_attributes_ = []
                    self._remote_attributes_ = []
                    
                    # Call the object's real __init__ method
                    Original.__init__(self, *args, **kwargs)

                @network_self.remote_procedure(authority)
                def removeFromNetwork(self):
                    # Remove the object from the network registry
                    if self._network_key_ != None:
                        network_self.remote_objects[self._network_key_] = None
                        network_self.remote_objects.pop(self._network_key_)
                        network_self.freeObjectKey(self._network_key_)
                        self._network_key_ = None
                    else:
                        raise Exception("Remote object has already been removed from the network.")

                #def deadreckon(self, attribute, extrapolation_function=(lambda dt, x, v: x+v*dt), tolerance = 0.1)
                #def synchronizeAttribute(self, attribute, rate_of_change, tolerance):
                #    if network_self.node_type == authority:
                #        self._remote_attributes_.append( Network.RemoteAttribute(attribute, rate_of_change, tolerance) )
                
                def synchAttribute(self, attribute, when, tolerance=0.0):
                    """Tells the network to synchronise the given attribute name

                       Note: The following may be rubbish.
                       Note: Attributes synchd on replicate will only be synched when
                              replicating a remote object to a non-auth who connected
                              after the remote object was created. Remote objects
                              created after a non-authority is connected will simply
                              have their __init__ method called as in regular
                              construction.
                    """
                    if network_self.node_type == authority:
                        if when == Network.SYNCH_CONTINUOUSLY:
                            print "Warning: Continuous synching of attributes is not currently implemented."
                        elif when == Network.SYNCH_ON_REPLICATE_ONLY:
                            self._synched_on_replicate_attributes_.append(attribute)
                    
            # Register the type as a remote type
            type_key = Original.__module__+'.'+Original.__name__
            network_self.remote_objects[type_key] = Replacement
            Replacement._network_type_key_ = type_key
            # Shhh, no-one will ever know...
            Replacement.__name__ = Original.__name__
            # TODO: Work out how to do this. It won't let me!
            #Replacement.__doc__  = Original.__doc__
            
            print "Registering remote object", Replacement, "with type key", type_key
            
            return Replacement
        return CLASS_DECORATOR
    
    class RemoteAttribute(object):
        def __init__(self, attribute, rate_of_change, tolerance):
            self.attribute = attribute
            self.rate_of_change = rate_of_change
            self.tolerance = tolerance
    
    def send(self, data, priority, reliability, order_channel, address=PyRakNet.SystemAddress.UNASSIGNED, broadcast=True):
        self.node.Send(data, priority, reliability, order_channel, address, broadcast)
